{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install -qU \\\n",
        "    datasets==2.19.1 \\\n",
        "    llama-index-core==0.11.9 \\\n",
        "    llama-index-llms-openai \\\n",
        "    llama-index-utils-workflow==0.2.1 \\\n",
        "    serpapi==0.1.5 \\\n",
        "    google-search-results==2.4.2 \\\n",
        "    semantic-router[pinecone]==0.0.65"
      ],
      "metadata": {
        "id": "WThmiBfvW7o0"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Knowledge Base Setup"
      ],
      "metadata": {
        "id": "RSNvnnmjhEZx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We'll be running our agent against a knowledge base \u2014 which requires a Pinecone index to be built.\n",
        "\n",
        "You can, if needed, skip this step and replace the `search` tool with a placeholder value if wanting to quickly test the structure of the Llama Index Workflow.\n",
        "\n",
        "If you want full functionality here, you do need to run this section \u2014 but we'll make it quick."
      ],
      "metadata": {
        "id": "yZCzed9thKZm"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Download a Dataset"
      ],
      "metadata": {
        "id": "EH0zitbAhI9N"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The first thing we need for an agent using RAG is somewhere we want to pull knowledge from. We will use v2 of the AI ArXiv dataset, available on Hugging Face Datasets at `jamescalam/ai-arxiv2-chunks`.\n",
        "\n",
        "Note: we're using the prechunked dataset. For the raw version see `jamescalam/ai-arxiv2`."
      ],
      "metadata": {
        "id": "h5UZcVZ1hO-O"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "dataset = load_dataset(\"jamescalam/ai-arxiv2-semantic-chunks\", split=\"train\")\n",
        "dataset"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_LTq8wOWhesx",
        "outputId": "53a018c6-d262-413d-c4e8-8836eefe1759"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:89: UserWarning: \n",
            "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
            "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
            "You will be able to reuse this secret in all of your notebooks.\n",
            "Please note that authentication is recommended but still optional to access public models or datasets.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Dataset({\n",
              "    features: ['id', 'title', 'content', 'prechunk_id', 'postchunk_id', 'arxiv_id', 'references'],\n",
              "    num_rows: 209760\n",
              "})"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dataset[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mqyP3a0_hqPm",
        "outputId": "cc1acf23-1781-4488-b086-7bef5d452beb"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'id': '2401.04088#0',\n",
              " 'title': 'Mixtral of Experts',\n",
              " 'content': '4 2 0 2 n a J 8 ] G L . s c [ 1 v 8 8 0 4 0 . 1 0 4 2 : v i X r a # Mixtral of Experts Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lample, L\u00c3\u00a9lio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Th\u00c3\u00a9ophile Gervet, Thibaut Lavril, Thomas Wang, Timoth\u00c3\u00a9e Lacroix, William El Sayed Abstract We introduce Mixtral 8x7B, a Sparse Mixture of Experts (SMoE) language model. Mixtral has the same architecture as Mistral 7B, with the difference that each layer is composed of 8 feedforward blocks (i.e. experts). For every token, at each layer, a router network selects two experts to process the current state and combine their outputs. Even though each token only sees two experts, the selected experts can be different at each timestep. As a result, each token has access to 47B parameters, but only uses 13B active parameters during inference. Mixtral was trained with a context size of 32k tokens and it outperforms or matches Llama 2 70B and GPT-3.5 across all evaluated benchmarks. In particular, Mixtral vastly outperforms Llama 2 70B on mathematics, code generation, and multilingual benchmarks. We also provide a model fine- tuned to follow instructions, Mixtral 8x7B \u00e2 Instruct, that surpasses GPT-3.5 Turbo, Claude-2.1, Gemini Pro, and Llama 2 70B \u00e2 chat model on human bench- marks. Both the base and instruct models are released under the Apache 2.0 license.',\n",
              " 'prechunk_id': '',\n",
              " 'postchunk_id': '2401.04088#1',\n",
              " 'arxiv_id': '2401.04088',\n",
              " 'references': ['1905.07830']}"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Construct Knowledge Base"
      ],
      "metadata": {
        "id": "TO1-zComhw3J"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "First initialize encoder model, for this we will need an [OpenAI API key](https://platform.openai.com/api-keys)."
      ],
      "metadata": {
        "id": "Cztd2YZ8h5VQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from getpass import getpass\n",
        "from semantic_router.encoders import OpenAIEncoder\n",
        "\n",
        "os.environ[\"OPENAI_API_KEY\"] = os.getenv(\"OPENAI_API_KEY\") or getpass(\"OpenAI API key: \")\n",
        "\n",
        "encoder = OpenAIEncoder(name=\"text-embedding-3-small\")"
      ],
      "metadata": {
        "id": "ZU92MBJxiUwq",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "96763a73-c3ff-4036-c33f-55222295a373"
      },
      "execution_count": 4,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "OpenAI API key: \u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Initialize our connection to Pinecone:"
      ],
      "metadata": {
        "id": "mpMLbVDLi1pC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pinecone import Pinecone\n",
        "\n",
        "# initialize connection to pinecone (get API key at app.pinecone.io)\n",
        "api_key = os.getenv(\"PINECONE_API_KEY\") or getpass(\"Pinecone API key: \")\n",
        "\n",
        "# configure client\n",
        "pc = Pinecone(api_key=api_key)"
      ],
      "metadata": {
        "id": "LzsjnjuCi3iy",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "19f0c0ca-3e38-4428-8b67-c2f4bbc46051"
      },
      "execution_count": 5,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Pinecone API key: \u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from pinecone import ServerlessSpec\n",
        "\n",
        "spec = ServerlessSpec(\n",
        "    cloud=\"aws\", region=\"us-east-1\"  # us-east-1\n",
        ")"
      ],
      "metadata": {
        "id": "GS535EvejxFM"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dims = len(encoder([\"some random text\"])[0])\n",
        "dims"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gFamwXVekeUx",
        "outputId": "d94021cc-c901-4d8c-ec94-e0b0bbd2dfc6"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "1536"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import time\n",
        "\n",
        "index_name = \"gpt-4o-research-agent\"\n",
        "\n",
        "# check if index already exists (it shouldn't if this is first time)\n",
        "if index_name not in pc.list_indexes().names():\n",
        "    # if does not exist, create index\n",
        "    pc.create_index(\n",
        "        index_name,\n",
        "        dimension=dims,  # dimensionality of embed 3\n",
        "        metric='dotproduct',\n",
        "        spec=spec\n",
        "    )\n",
        "    # wait for index to be initialized\n",
        "    while not pc.describe_index(index_name).status['ready']:\n",
        "        time.sleep(1)\n",
        "\n",
        "# connect to index\n",
        "index = pc.Index(index_name)\n",
        "time.sleep(1)\n",
        "# view index stats\n",
        "index.describe_index_stats()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YBzhQFERkjwN",
        "outputId": "1a7e9832-938e-4675-a7bd-fc9030a90cef"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'dimension': 1536,\n",
              " 'index_fullness': 0.0,\n",
              " 'namespaces': {'': {'vector_count': 10000}},\n",
              " 'total_vector_count': 10000}"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Populate the knowledge base:"
      ],
      "metadata": {
        "id": "SEP_TzsMkldk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from tqdm.auto import tqdm\n",
        "\n",
        "# easier to work with dataset as pandas dataframe\n",
        "data = dataset.to_pandas().iloc[:10000]\n",
        "\n",
        "batch_size = 128\n",
        "\n",
        "for i in tqdm(range(0, len(data), batch_size)):\n",
        "    i_end = min(len(data), i+batch_size)\n",
        "    batch = data[i:i_end].to_dict(orient=\"records\")\n",
        "    # get batch of data\n",
        "    metadata = [{\n",
        "        \"title\": r[\"title\"],\n",
        "        \"content\": r[\"content\"],\n",
        "        \"arxiv_id\": r[\"arxiv_id\"],\n",
        "        \"references\": r[\"references\"].tolist()\n",
        "    } for r in batch]\n",
        "    # generate unique ids for each chunk\n",
        "    ids = [r[\"id\"] for r in batch]\n",
        "    # get text content to embed\n",
        "    content = [r[\"content\"] for r in batch]\n",
        "    # embed text\n",
        "    embeds = encoder(content)\n",
        "    # add to Pinecone\n",
        "    index.upsert(vectors=zip(ids, embeds, metadata))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "58f0cf25a2ec4aa3a0dabcb8a2a7c36f",
            "7f441d20804b4b829d7f09484c0369cb",
            "d8a83df6fd834d128e96057e664d9162",
            "288972f4d50549edbe00ee46407fef28",
            "e3d4bdfe01214bd6a19775f15b58f651",
            "8c46848d4c7949e1abe0c8d430f67da0",
            "427c3276bb9a421eb9814711a4934876",
            "d14b8a66298b477ca6b070d99ed86a50",
            "f641d0870d9e4407b185f622c839bf21",
            "8167b2cf4ee04576bd7f5e232be09280",
            "97d3249e633b4cc6b5beb7e51c6aefdb"
          ]
        },
        "id": "SYdvU-fjkold",
        "outputId": "9f26cf1b-7c56-407a-ff7b-e1757f5e6c4b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/79 [00:02<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "58f0cf25a2ec4aa3a0dabcb8a2a7c36f"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Agent Components"
      ],
      "metadata": {
        "id": "FUh_C_5HhYh6"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Tools"
      ],
      "metadata": {
        "id": "0RS6Ow8FLmZE"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We define the separate tool functions. When integrating with our graph all of these will be executed using the same `run_tools` class - which we will define later.\n",
        "\n",
        "For now, let's define the functions that our agent will have access to."
      ],
      "metadata": {
        "id": "51qdvTLIYKL6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import requests\n",
        "import re\n",
        "\n",
        "# our regex\n",
        "abstract_pattern = re.compile(\n",
        "    r'\\s*Abstract:\\s*(.*?)\\s*',\n",
        "    re.DOTALL\n",
        ")\n",
        "\n",
        "async def fetch_arxiv(arxiv_id: str):\n",
        "    \"\"\"Gets the abstract from an ArXiv paper given the arxiv ID. Useful for\n",
        "    finding high-level context about a specific paper.\"\"\"\n",
        "    print(\">>> fetch_arxiv\")\n",
        "    # get paper page in html\n",
        "    res = requests.get(\n",
        "        f\"https://export.arxiv.org/abs/{arxiv_id}\"\n",
        "    )\n",
        "    # search html for abstract\n",
        "    re_match = abstract_pattern.search(res.text)\n",
        "    # return abstract text\n",
        "    return re_match.group(1)"
      ],
      "metadata": {
        "id": "Gxe_NBTHLm01"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Web Search"
      ],
      "metadata": {
        "id": "sXmNfv63L2Dw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The web search tool will provide the agent with access to web search. It will be instructed to use this for more general knowledge queries."
      ],
      "metadata": {
        "id": "bpRQvPiFLvF9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from serpapi import GoogleSearch\n",
        "from getpass import getpass\n",
        "import os\n",
        "\n",
        "serpapi_params = {\n",
        "    \"engine\": \"google\",\n",
        "    \"api_key\": os.getenv(\"SERPAPI_KEY\") or getpass(\"SerpAPI key: \")\n",
        "}\n",
        "\n",
        "async def web_search(query: str):\n",
        "    \"\"\"Finds general knowledge information using Google search. Can also be used\n",
        "    to augment more 'general' knowledge to a previous specialist query.\"\"\"\n",
        "    print(\">>> web_search\")\n",
        "    search = GoogleSearch({\n",
        "        **serpapi_params,\n",
        "        \"q\": query,\n",
        "        \"num\": 5\n",
        "    })\n",
        "    results = search.get_dict()[\"organic_results\"]\n",
        "    contexts = \"\\n---\\n\".join(\n",
        "        [\"\\n\".join([x[\"title\"], x[\"snippet\"], x[\"link\"]]) for x in results]\n",
        "    )\n",
        "    return contexts"
      ],
      "metadata": {
        "id": "ZRLxR_eKLvgz",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c7685b9a-e98d-4411-9ee3-c9ada913b1e4"
      },
      "execution_count": 10,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "SerpAPI key: \u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\u00b7\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Rag Search"
      ],
      "metadata": {
        "id": "NeckltzgL92n"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We provide two RAG-focused tools for our agent. The `rag_search` allows the agent to perform a simple RAG search for some information across all indexed research papers. The `rag_search_filter` also searches, but within a specific paper which is filtered for via the `arxiv_id` parameter.\n",
        "\n",
        "We also define the `format_rag_contexts` function to handle the transformation of our Pinecone results from a JSON object to a readble plaintext format."
      ],
      "metadata": {
        "id": "uEU4NDWlL-Ox"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def format_rag_contexts(matches: list):\n",
        "    contexts = []\n",
        "    for x in matches:\n",
        "        text = (\n",
        "            f\"Title: {x['metadata']['title']}\\n\"\n",
        "            f\"Content: {x['metadata']['content']}\\n\"\n",
        "            f\"ArXiv ID: {x['metadata']['arxiv_id']}\\n\"\n",
        "            f\"Related Papers: {x['metadata']['references']}\\n\"\n",
        "        )\n",
        "        contexts.append(text)\n",
        "    context_str = \"\\n---\\n\".join(contexts)\n",
        "    return context_str\n",
        "\n",
        "async def rag_search_filter(query: str, arxiv_id: str):\n",
        "    \"\"\"Finds information from our ArXiv database using a natural language query\n",
        "    and a specific ArXiv ID. Allows us to learn more details about a specific paper.\"\"\"\n",
        "    print(\">>> rag_search_filter\")\n",
        "    xq = await encoder.acall([query])\n",
        "    xc = index.query(vector=xq, top_k=6, include_metadata=True, filter={\"arxiv_id\": arxiv_id})\n",
        "    context_str = format_rag_contexts(xc[\"matches\"])\n",
        "    return context_str\n",
        "\n",
        "async def rag_search(query: str):\n",
        "    \"\"\"Finds specialist information on AI using a natural language query.\"\"\"\n",
        "    print(\">>> rag_search\")\n",
        "    xq = await encoder.acall([query])\n",
        "    xc = index.query(vector=xq, top_k=2, include_metadata=True)\n",
        "    context_str = format_rag_contexts(xc[\"matches\"])\n",
        "    return context_str"
      ],
      "metadata": {
        "id": "kuv6bK99L-no"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "async def final_answer(\n",
        "    introduction: str,\n",
        "    research_steps: str,\n",
        "    main_body: str,\n",
        "    conclusion: str,\n",
        "    sources: str\n",
        "):\n",
        "    \"\"\"Returns a natural language response to the user in the form of a research\n",
        "    report. There are several sections to this report, those are:\n",
        "    - `introduction`: a short paragraph introducing the user's question and the\n",
        "    topic we are researching.\n",
        "    - `research_steps`: a few bullet points explaining the steps that were taken\n",
        "    to research your report.\n",
        "    - `main_body`: this is where the bulk of high quality and concise\n",
        "    information that answers the user's question belongs. It is 3-4 paragraphs\n",
        "    long in length.\n",
        "    - `conclusion`: this is a short single paragraph conclusion providing a\n",
        "    concise but sophisticated view on what was found.\n",
        "    - `sources`: a bulletpoint list provided detailed sources for all information\n",
        "    referenced during the research process\n",
        "    \"\"\"\n",
        "    print(\">>> final_answer\")\n",
        "    if type(research_steps) is list:\n",
        "        research_steps = \"\\n\".join([f\"- {r}\" for r in research_steps])\n",
        "    if type(sources) is list:\n",
        "        sources = \"\\n\".join([f\"- {s}\" for s in sources])\n",
        "    return \"\""
      ],
      "metadata": {
        "id": "OGsCy9ZMVJmy"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Oracle LLM"
      ],
      "metadata": {
        "id": "WUm29WrkVNYh"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Our prompt for the Oracle will emphasize it's decision making ability within the `system_prompt`, leave a placeholder for us to later insert `chat_history`, and provide a place for us to insert the user `input`."
      ],
      "metadata": {
        "id": "sqM2mkxsVOoG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "system_prompt = \"\"\"You are the oracle, the great AI decision maker.\n",
        "Given the user's query you must decide what to do with it based on the\n",
        "list of tools provided to you.\n",
        "\n",
        "If you see that a tool has been used (in the scratchpad) with a particular\n",
        "query, do NOT use that same tool with the same query again. Also, do NOT use\n",
        "any tool more than twice (ie, if the tool appears in the scratchpad twice, do\n",
        "not use it again).\n",
        "\n",
        "You should aim to collect information from a diverse range of sources before\n",
        "providing the answer to the user. Once you have collected plenty of information\n",
        "to answer the user's question (stored in the scratchpad) use the final_answer\n",
        "tool.\"\"\""
      ],
      "metadata": {
        "id": "AW96-NkUVN5u"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The oracle agent will be provided the tools we previously built."
      ],
      "metadata": {
        "id": "9L59sEakdMd2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.core.tools import FunctionTool\n",
        "\n",
        "tools = [\n",
        "    FunctionTool.from_defaults(async_fn=fetch_arxiv),\n",
        "    FunctionTool.from_defaults(async_fn=web_search),\n",
        "    FunctionTool.from_defaults(async_fn=rag_search_filter),\n",
        "    FunctionTool.from_defaults(async_fn=rag_search),\n",
        "    FunctionTool.from_defaults(async_fn=final_answer),\n",
        "]"
      ],
      "metadata": {
        "id": "sELUig_VdXK7"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from llama_index.llms.openai import OpenAI\n",
        "\n",
        "llm = OpenAI(\n",
        "    model=\"gpt-4o\",\n",
        "    additional_kwargs={\"tool_choice\": \"required\"}\n",
        ")"
      ],
      "metadata": {
        "id": "LNiTGVrMdtN4"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Events"
      ],
      "metadata": {
        "id": "kcXeJrFZa9eG"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We need to create a few events for our workflow. Llama-index comes with a few predefined event types, two of which we will use (`StartEvent` and `StopEvent`). However, we need to define a few additional custom event types - these are:\n",
        "\n",
        "* `InputEvent` to handle new messages and prepare chat history.\n",
        "\n",
        "* `ToolCallEvent` to trigger tool calls."
      ],
      "metadata": {
        "id": "RzCsuWEsbAuS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.core.llms import ChatMessage\n",
        "from llama_index.core.tools import ToolSelection, ToolOutput\n",
        "from llama_index.core.workflow import Event\n",
        "\n",
        "\n",
        "class InputEvent(Event):\n",
        "    input: list[ChatMessage]\n",
        "\n",
        "\n",
        "class ToolCallEvent(Event):\n",
        "    id: str\n",
        "    name: str\n",
        "    params: dict"
      ],
      "metadata": {
        "id": "FXdBlptmbAQk"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we build the workflow. Workflows consist of a single `Workflow` class with multiple `steps`. Each step is like a compute/execution step in our agentic flow.\n",
        "\n",
        "We control which step is triggered by using different `Event` types. Each step consumes a different type of event, like `InputEvent` or `ToolCallEvent`. Additionally, our workflow begins and ends with the `StartEvent` and `StopEvent` events respectively."
      ],
      "metadata": {
        "id": "k9aE7MiznZCj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.core.workflow import Workflow, StartEvent, StopEvent, step\n",
        "from llama_index.core.memory import ChatMemoryBuffer\n",
        "from llama_index.core.llms import MessageRole\n",
        "\n",
        "class ResearchAgent(Workflow):\n",
        "    def __init__(\n",
        "        self,\n",
        "        *args: any,\n",
        "        oracle: OpenAI,\n",
        "        tools: list[FunctionTool],\n",
        "        timeout: int = 20,\n",
        "    ):\n",
        "        super().__init__(*args)\n",
        "        self._timeout = timeout\n",
        "        self.oracle = oracle\n",
        "        self.tools = tools\n",
        "        self.get_tool = {tool.metadata.get_name(): tool for tool in self.tools}\n",
        "        # initialize chat history/memory with system prompt\n",
        "        self.sys_msg = ChatMessage(\n",
        "            role=MessageRole.SYSTEM,\n",
        "            content=system_prompt\n",
        "        )\n",
        "        self.memory = ChatMemoryBuffer.from_defaults(llm=OpenAI())\n",
        "\n",
        "    @step\n",
        "    async def prepare_chat_history(self, ev: StartEvent) -> InputEvent:\n",
        "        # clear memory\n",
        "        self.memory = ChatMemoryBuffer.from_defaults(llm=OpenAI())\n",
        "        self.memory.put(message=self.sys_msg)\n",
        "        # get user input\n",
        "        user_input = ev.input\n",
        "        user_msg = ChatMessage(role=\"user\", content=user_input)\n",
        "        self.memory.put(message=user_msg)\n",
        "        # get chat history\n",
        "        chat_history = self.memory.get()\n",
        "        # return input event\n",
        "        return InputEvent(input=chat_history)\n",
        "\n",
        "    @step\n",
        "    async def handle_llm_input(self, ev: InputEvent) -> ToolCallEvent | StopEvent:\n",
        "        chat_history = ev.input\n",
        "        # get oracle response\n",
        "        response = await self.oracle.achat_with_tools(\n",
        "            tools=self.tools,\n",
        "            chat_history=chat_history,\n",
        "            tool_choice=\"required\",\n",
        "        )\n",
        "        # add response to chat history / memory\n",
        "        self.memory.put(message=response.message)\n",
        "        # get tool calls\n",
        "        tool_calls = self.oracle.get_tool_calls_from_response(\n",
        "            response\n",
        "        )\n",
        "        # if final_answer tool used we return to the user with the StopEvent\n",
        "        if tool_calls[-1].tool_name == \"final_answer\":\n",
        "            return StopEvent(result={\"response\": tool_calls[-1].tool_kwargs})\n",
        "        else:\n",
        "            # return tool call event\n",
        "            return ToolCallEvent(\n",
        "                id=tool_calls[-1].tool_id,\n",
        "                name=tool_calls[-1].tool_name,\n",
        "                params=tool_calls[-1].tool_kwargs,\n",
        "            )\n",
        "\n",
        "    @step\n",
        "    async def run_tool(self, ev: ToolCallEvent) -> InputEvent:\n",
        "        tool_name = ev.name\n",
        "        additional_kwargs = {\n",
        "            \"tool_call_id\": ev.id,\n",
        "            \"name\": tool_name\n",
        "        }\n",
        "        # get chosen tool\n",
        "        tool = self.get_tool.get(tool_name)\n",
        "        if not tool:\n",
        "            tool_msg = ChatMessage(\n",
        "                role=\"tool\",\n",
        "                content=f\"Tool {tool_name} not found\",\n",
        "                additional_kwargs=additional_kwargs\n",
        "            )\n",
        "        else:\n",
        "            # now call tool\n",
        "            tool_output = await tool.acall(**ev.params)\n",
        "            tool_msg = ChatMessage(\n",
        "                role=\"tool\",\n",
        "                content=tool_output.content,\n",
        "                additional_kwargs=additional_kwargs\n",
        "            )\n",
        "        self.memory.put(message=tool_msg)\n",
        "        chat_history = self.memory.get()\n",
        "        return InputEvent(input=chat_history)\n",
        "\n"
      ],
      "metadata": {
        "id": "qXq1QbvScZeZ"
      },
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.core.workflow import draw_all_possible_flows\n",
        "\n",
        "draw_all_possible_flows(ResearchAgent, filename=\"research_agent.html\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pXs5pRIg4El9",
        "outputId": "3e5ff233-83e0-4345-c63b-e72753a916b7"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "research_agent.html\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-39-75edaaa7c19e>:3: DeprecationWarning: Call to deprecated function (or staticmethod) draw_all_possible_flows. (Install `llama-index-utils-workflow` and use the import `from llama_index.utils.workflow` instead.)\n",
            "  draw_all_possible_flows(ResearchAgent, filename=\"research_agent.html\")\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Initialize the workflow:"
      ],
      "metadata": {
        "id": "cO99bN2pnLOB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "agent = ResearchAgent(\n",
        "    oracle=llm,\n",
        "    tools=tools,\n",
        "    timeout=30,\n",
        ")"
      ],
      "metadata": {
        "id": "gVYXBk6SnEYE"
      },
      "execution_count": 46,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "res = await agent.run(input=\"tell me about AI\")\n",
        "res[\"response\"]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2ynhR2BTqRim",
        "outputId": "30c3a385-57d6-4136-cbdf-1e51d5e06466"
      },
      "execution_count": 47,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            ">>> web_search\n",
            ">>> rag_search\n",
            ">>> fetch_arxiv\n",
            ">>> rag_search\n",
            ">>> web_search\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'introduction': 'Artificial Intelligence (AI) is a transformative technology that has evolved significantly since its inception. It encompasses a wide range of techniques and applications aimed at replicating human intelligence in machines.',\n",
              " 'research_steps': '1. Conducted a web search to gather general information about AI.\\n2. Utilized a specialized AI research tool to obtain detailed insights from academic papers.\\n3. Retrieved specific historical context and developments in AI from both general and specialized sources.',\n",
              " 'main_body': \"Artificial Intelligence (AI) refers to the simulation of human intelligence processes by machines, particularly computer systems. These processes include learning (the acquisition of information and rules for using the information), reasoning (using rules to reach approximate or definite conclusions), and self-correction. Specific applications of AI include expert systems, natural language processing (NLP), speech recognition, and machine vision.\\n\\nThe concept of AI dates back to ancient times, with myths and stories about artificial beings endowed with intelligence or consciousness. However, the formal foundation of AI as a field of study was laid in the mid-20th century. British logician and computer scientist Alan Turing is often credited with pioneering AI through his development of the Turing Test, which assesses a machine's ability to exhibit intelligent behavior equivalent to, or indistinguishable from, that of a human.\\n\\nIn the 1950s and 1960s, AI research focused on symbolic methods and problem-solving. John McCarthy, Marvin Minsky, Allen Newell, and Herbert A. Simon were among the key figures who contributed to the development of AI during this period. The term 'Artificial Intelligence' was coined by McCarthy in 1956 during the Dartmouth Conference, which is considered the birth of AI as an academic discipline.\\n\\nThe field has since evolved through various phases, including the rise of machine learning and neural networks in the 1980s and 1990s. The advent of deep learning and reinforcement learning in the 2010s marked a significant shift, enabling the development of AI systems capable of outperforming humans in specific tasks, such as playing complex games like Go and chess. Today, AI continues to advance, driven by large-scale data, powerful computational resources, and innovative algorithms.\",\n",
              " 'conclusion': 'AI has come a long way from its early conceptual stages to becoming a cornerstone of modern technology. Its continuous evolution promises to bring even more sophisticated and beneficial applications, transforming various aspects of human life and industry.',\n",
              " 'sources': '- Google Cloud: What is Artificial Intelligence?\\n- IBM: What is Artificial Intelligence?\\n- TechTarget: Artificial Intelligence Explained\\n- ArXiv: The Rise and Potential of Large Language Model Based Agents: A Survey\\n- Wikipedia: History of Artificial Intelligence\\n- Harvard University: The History of Artificial Intelligence'}"
            ]
          },
          "metadata": {},
          "execution_count": 47
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "---"
      ],
      "metadata": {
        "id": "Looyvi2I4gvc"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's test with async. To avoid overwriting state with asynchronous runs taking place we will initialize three new agent instances."
      ],
      "metadata": {
        "id": "hBT8Wtgu6ax9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "agent1 = ResearchAgent(oracle=llm, tools=tools, timeout=30)\n",
        "agent2 = ResearchAgent(oracle=llm, tools=tools, timeout=30)\n",
        "agent3 = ResearchAgent(oracle=llm, tools=tools, timeout=30)"
      ],
      "metadata": {
        "id": "6qeFbqCQ8RJ3"
      },
      "execution_count": 48,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import asyncio\n",
        "\n",
        "calls = [\n",
        "    agent.run(input=\"tell me about AI\"),\n",
        "    agent1.run(input=\"tell me about AI\"),\n",
        "    agent2.run(input=\"what is RAG?\"),\n",
        "    agent3.run(input=\"what is the latest LLM from OpenAI?\")\n",
        "]\n",
        "\n",
        "outputs = await asyncio.gather(*calls)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BgwZddGK6PaX",
        "outputId": "7992d327-b8d5-46f5-cf77-da444e46f6b6"
      },
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            ">>> web_search\n",
            ">>> web_search\n",
            ">>> web_search\n",
            ">>> web_search\n",
            ">>> rag_search\n",
            ">>> rag_search\n",
            ">>> rag_search\n",
            ">>> fetch_arxiv\n",
            ">>> fetch_arxiv\n",
            ">>> fetch_arxiv\n",
            ">>> rag_search\n",
            ">>> fetch_arxiv\n",
            ">>> rag_search_filter\n",
            ">>> rag_search_filter\n",
            ">>> rag_search_filter\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "outputs[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AYCUOBok66YU",
        "outputId": "59d9a095-71b6-4641-d894-eb7875633360"
      },
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'response': {'introduction': 'Artificial Intelligence (AI) continues to evolve rapidly, with 2023 witnessing significant advancements and trends that are shaping the future of technology. This report delves into the latest AI research trends of 2023, focusing on the rise of generative AI, the development of autonomous cognitive entities, and other key areas.',\n",
              "  'research_steps': '- Conducted a web search to identify the latest AI research trends in 2023.\\n- Utilized RAG search to gather specialist information on AI trends.\\n- Retrieved abstracts and detailed information from specific ArXiv papers related to generative AI and autonomous cognitive entities.\\n- Analyzed and synthesized the collected information to provide a comprehensive overview.',\n",
              "  'main_body': 'In 2023, the field of AI has been marked by several groundbreaking trends, with generative AI taking center stage. Generative AI, which involves creating new content from existing data, has seen explosive growth. Tools like ChatGPT and Claude have demonstrated the potential of large language models (LLMs) to generate human-like text, leading to widespread adoption in various industries. These models are not only enhancing user interactions but also driving innovations in content creation, customer service, and more.\\n\\nAnother significant trend is the development of autonomous cognitive entities. The Conceptual Framework for Autonomous Cognitive Entities (ACE) introduces a novel cognitive architecture that enables machines and software agents to operate more independently. This framework leverages the capabilities of the latest generative AI technologies, including LLMs and multimodal generative models (MMMs), to build autonomous, agentic systems. The ACE framework comprises six layers, each playing a distinct role in setting moral compasses, strategic thinking, task selection, and execution. This development is poised to revolutionize the way autonomous systems are designed and implemented.\\n\\nAdditionally, the integration of AI in various sectors continues to expand. AI-powered healthcare solutions, computer vision advancements, and AI-driven cybersecurity measures are becoming increasingly prevalent. These applications highlight the versatility of AI technologies and their potential to address complex challenges across different domains.\\n\\nThe research also underscores the importance of ethical considerations and safety in AI development. As AI systems become more autonomous and capable, ensuring their alignment with human values and societal norms is crucial. Researchers are actively exploring ways to embed ethical principles into AI frameworks to mitigate potential risks and enhance the trustworthiness of AI systems.',\n",
              "  'conclusion': 'The AI landscape in 2023 is characterized by remarkable advancements in generative AI and the emergence of autonomous cognitive entities. These trends are driving innovation across various industries and reshaping the future of technology. As AI continues to evolve, addressing ethical considerations and ensuring the safe deployment of AI systems will be paramount.',\n",
              "  'sources': \"- McKinsey Global Survey on AI, 2023\\n- Technology Review, 'Four trends that changed AI in 2023'\\n- ArXiv papers: 'SAPIEN: Affective Virtual Agents Powered by Large Language Models' (2308.03022), 'Conceptual Framework for Autonomous Cognitive Entities' (2310.06775)\"}}"
            ]
          },
          "metadata": {},
          "execution_count": 50
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "outputs[1]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tNw2XibX8v5p",
        "outputId": "fe20f2a6-9404-47b6-e70b-016127179038"
      },
      "execution_count": 51,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'response': {'introduction': 'Artificial Intelligence (AI) is a transformative field of computer science focused on creating systems capable of performing tasks that typically require human intelligence. This includes activities such as learning, reasoning, problem-solving, perception, and language understanding.',\n",
              "  'research_steps': '1. Conducted a web search to gather general information about AI.\\n2. Performed a specialized search in AI literature to obtain detailed insights.\\n3. Retrieved and reviewed abstracts and content from relevant ArXiv papers.\\n4. Synthesized information from multiple sources to provide a comprehensive overview.',\n",
              "  'main_body': \"Artificial Intelligence (AI) is a broad and dynamic field within computer science that aims to create machines capable of intelligent behavior. The concept of AI dates back to ancient philosophical discussions about the nature of intelligence and the possibility of non-human entities possessing it. In the 1950s, Alan Turing formalized the idea with the Turing Test, which evaluates a machine's ability to exhibit intelligent behavior indistinguishable from that of a human.\\n\\nAI systems are often referred to as 'agents,' which are entities capable of perceiving their environment through sensors, making decisions, and taking actions using actuators. These agents can range from simple rule-based systems to complex neural networks that learn and adapt over time. The development of AI has gone through several stages, from symbolic AI, which relies on explicit rules and logic, to modern approaches like machine learning and deep learning, which use vast amounts of data to train models.\\n\\nOne of the most significant advancements in AI is the development of large language models, which can understand and generate human-like text. These models are based on neural networks and have shown remarkable capabilities in tasks such as translation, summarization, and even creative writing. The potential applications of AI are vast, spanning industries such as healthcare, finance, transportation, and entertainment.\\n\\nDespite its rapid progress, AI also raises important ethical and societal questions. Issues such as bias in AI algorithms, the impact on employment, and the need for transparency and accountability in AI systems are critical areas of ongoing research and debate.\",\n",
              "  'conclusion': 'AI is a rapidly evolving field that holds great promise for transforming various aspects of society. While it offers numerous benefits, it also presents challenges that need to be addressed through careful research and thoughtful implementation. As AI continues to advance, it will be essential to balance innovation with ethical considerations to ensure its positive impact on the world.',\n",
              "  'sources': '- Google Cloud: What is Artificial Intelligence?\\n- IBM: What is Artificial Intelligence?\\n- TechTarget: Artificial Intelligence Explained\\n- ArXiv Paper: The Rise and Potential of Large Language Model Based Agents: A Survey (2309.07864)'}}"
            ]
          },
          "metadata": {},
          "execution_count": 51
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "outputs[2]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "O-4-rjVz87IS",
        "outputId": "eaf90697-b6a2-458b-84bd-e5299658c801"
      },
      "execution_count": 52,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'response': {'introduction': 'Retrieval-Augmented Generation (RAG) is an advanced AI framework that enhances the capabilities of generative models by integrating external knowledge sources. This approach aims to improve the accuracy and reliability of AI-generated content by leveraging both the inherent knowledge of the language model and additional information retrieved from external databases.',\n",
              "  'research_steps': '- Conducted a web search to gather general information about RAG.\\n- Performed a specialized RAG search to obtain detailed insights.\\n- Retrieved and analyzed an ArXiv paper (ID: 2308.03983) to understand the technical aspects of RAG.\\n- Filtered the ArXiv paper for specific details about RAG.',\n",
              "  'main_body': 'Retrieval-Augmented Generation (RAG) is a technique designed to enhance the performance of generative AI models by incorporating external knowledge sources. Traditional generative models rely solely on their pre-trained knowledge, which can sometimes lead to inaccuracies or hallucinations. RAG addresses this limitation by integrating a retrieval mechanism that fetches relevant information from external databases or knowledge bases, thereby augmenting the model\\'s responses with up-to-date and contextually relevant data.\\n\\nThe core idea behind RAG is to allow the language model to access and utilize external information dynamically. This is achieved through a two-step process: retrieval and generation. In the retrieval step, the model identifies and retrieves pertinent information from an external source based on the input query. In the generation step, the model combines this retrieved information with its inherent knowledge to produce a more accurate and contextually appropriate response.\\n\\nA study detailed in the ArXiv paper \"SimplyRetrieve: A Private and Lightweight Retrieval-Centric Generative AI Tool\" (ID: 2308.03983) compares RAG with other approaches like Retrieval-Centric Generation (RCG) and Retrieval-OFF Generation (ROG). The findings indicate that RAG allows for a more flexible integration of the language model\\'s inherent knowledge with externally retrieved data. This flexibility can sometimes lead to partially erroneous information, but it generally enhances the model\\'s ability to provide accurate and relevant responses. The study also highlights the importance of prompt engineering in optimizing the performance of RAG-based systems.\\n\\nRAG\\'s ability to dynamically incorporate external knowledge makes it particularly useful in applications requiring up-to-date information, such as question answering, customer support, and content generation. By leveraging external databases, RAG can provide more reliable and contextually accurate responses, thereby improving the overall user experience.',\n",
              "  'conclusion': 'Retrieval-Augmented Generation (RAG) represents a significant advancement in the field of generative AI. By integrating external knowledge sources, RAG enhances the accuracy and reliability of AI-generated content. While it has its challenges, such as the potential for partially erroneous information, the benefits of improved contextual relevance and up-to-date responses make RAG a valuable tool in various AI applications.',\n",
              "  'sources': '- AWS: What is Retrieval-Augmented Generation?\\n- Google Cloud: What Is Retrieval Augmented Generation (RAG)?\\n- NVIDIA Blog: What Is Retrieval-Augmented Generation aka RAG\\n- ArXiv Paper: SimplyRetrieve: A Private and Lightweight Retrieval-Centric Generative AI Tool (ID: 2308.03983)'}}"
            ]
          },
          "metadata": {},
          "execution_count": 52
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "outputs[3]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3DStj2Zu89ad",
        "outputId": "3c9c0aff-7c81-4938-88c9-49362b864ac9"
      },
      "execution_count": 53,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'response': {'introduction': \"OpenAI has been at the forefront of developing advanced large language models (LLMs) that push the boundaries of artificial intelligence. The latest model from OpenAI, as of 2023, is known as 'o1'.\",\n",
              "  'research_steps': \"1. Conducted a web search to find the latest information on OpenAI's LLMs.\\n2. Reviewed multiple sources to confirm the details about the latest model.\\n3. Compiled the information into a concise report.\",\n",
              "  'main_body': \"OpenAI's latest large language model, introduced in 2023, is called 'o1'. This model represents a significant advancement in the field of AI, particularly in its ability to perform complex reasoning tasks. The 'o1' model has been trained using reinforcement learning techniques, which enable it to think before responding, thereby improving its inferencing capabilities. This makes 'o1' particularly adept at handling complex, multi-step tasks that require a higher level of cognitive processing compared to its predecessors.\\n\\nIn addition to 'o1', OpenAI has also introduced another model named 'GPT-4o'. This model is designed to be a high-intelligence flagship model for complex tasks, offering improvements in speed and cost-efficiency over previous models like GPT-4 Turbo. 'GPT-4o' can handle up to 128,000 tokens, making it suitable for extensive and intricate tasks.\\n\\nThese advancements highlight OpenAI's commitment to pushing the boundaries of what is possible with LLMs. The introduction of 'o1' and 'GPT-4o' not only enhances the capabilities of AI but also sets new standards for future developments in the field.\",\n",
              "  'conclusion': \"OpenAI's latest LLMs, 'o1' and 'GPT-4o', represent significant strides in AI technology, particularly in complex reasoning and task handling. These models are set to redefine the capabilities and applications of AI in various domains.\",\n",
              "  'sources': \"- OpenAI official website\\n- Psychology Today article on OpenAI's latest LLM\\n- OpenAI API documentation\"}}"
            ]
          },
          "metadata": {},
          "execution_count": 53
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "---"
      ],
      "metadata": {
        "id": "Y-2YzSQA9A8L"
      }
    }
  ]
}